#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <netdb.h>
#include <sys/select.h>

#define MAX_BUFFER 1024
#define MAX_NAME 50
#define MAX_GROUP 20

// 消息类型定义
#define REGISTER 1
#define LOGIN 2
#define LOGOUT 3
#define LIST_USERS 4
#define LIST_GROUPS 5
#define PRIVATE_MSG 6
#define CREATE_GROUP 7
#define JOIN_GROUP 8
#define GROUP_MSG 9
#define LEAVE_GROUP 10
#define BROADCAST 11

typedef struct {
    int type;
    char sender[MAX_NAME];
    char target[MAX_NAME];
    char message[MAX_BUFFER];
} ChatMessage;

void print_menu() {
    printf("\n===== 聊天菜单 =====\n");
    printf("1. 获取在线用户列表\n");
    printf("2. 获取群组列表\n");
    printf("3. 发送私信\n");
    printf("4. 创建群组\n");
    printf("5. 加入群组\n");
    printf("6. 发送群组消息\n");
    printf("7. 离开群组\n");
    printf("8. 退出登录\n");
    printf("9. 广播消息\n");
    printf("====================\n");
    printf("输入 -help 获取此菜单\n");
}

void register_user(int udp_sock, struct sockaddr_in* server_addr) {
    char username[MAX_NAME], password[MAX_NAME];
    
    printf("===== 用户注册 =====\n");
    printf("用户名: ");
    fgets(username, MAX_NAME, stdin);
    username[strcspn(username, "\n")] = '\0';
    
    printf("密码: ");
    fgets(password, MAX_NAME, stdin);
    password[strcspn(password, "\n")] = '\0';

    ChatMessage reg_msg;
    reg_msg.type = REGISTER;
    strncpy(reg_msg.sender, username, MAX_NAME);
    strncpy(reg_msg.message, password, MAX_BUFFER);

    sendto(udp_sock, &reg_msg, sizeof(reg_msg), 0, 
           (struct sockaddr*)server_addr, sizeof(*server_addr));

    // 接收注册响应
    ChatMessage reg_resp;
    socklen_t server_len = sizeof(*server_addr);
    recvfrom(udp_sock, &reg_resp, sizeof(reg_resp), 0, 
             (struct sockaddr*)server_addr, &server_len);
    
    if (strcmp(reg_resp.message, "SUCCESS") != 0) {
        printf("注册失败: %s\n", reg_resp.message);
    } else {
        printf("注册成功！\n");
    }
}

void login_user(int tcp_sock, char* username) {
    ChatMessage login_msg;  
    login_msg.type = LOGIN;
    strncpy(login_msg.sender, username, MAX_NAME);
    
    printf("密码: ");
    char password[MAX_NAME];
    fgets(password, MAX_NAME, stdin);
    password[strcspn(password, "\n")] = '\0';
    
    strncpy(login_msg.message, password, MAX_BUFFER);
    send(tcp_sock, &login_msg, sizeof(login_msg), 0);

    // 接收登录响应
    ChatMessage login_resp;
    recv(tcp_sock, &login_resp, sizeof(login_resp), 0);
    if (strcmp(login_resp.message, "SUCCESS") != 0) {
        printf("登录失败: %s\n", login_resp.message);
    } else {
        printf("登录成功! 欢迎 %s\n", username);
    }
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        fprintf(stderr, "用法: %s <服务器IP> <端口号>\n", argv[0]);
        exit(1);
    }

    char *server_ip = argv[1];
    int port = atoi(argv[2]);

    // 创建UDP套接字用于注册
    int udp_sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (udp_sock < 0) {
        perror("UDP套接字创建失败");
        exit(1);
    }

    // 创建TCP套接字用于登录后通信
    int tcp_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (tcp_sock < 0) {
        perror("TCP套接字创建失败");
        close(udp_sock);
        exit(1);
    }

    // 设置服务器地址
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    
    if (inet_pton(AF_INET, server_ip, &server_addr.sin_addr) <= 0) {
        perror("地址转换失败");
        close(udp_sock);
        close(tcp_sock);
        exit(1);
    }

    char username[MAX_NAME];  // 用于存储用户名
    int logged_in = 0; // 标识用户是否已登录

    while (1) {
        int option;
        printf("\n选择操作 (1: 注册; 2: 登录; 3: 退出): ");
        scanf("%d", &option);
        getchar(); // 清除换行符
        
        if (option == 1) { // 注册
            register_user(udp_sock, &server_addr);
        } else if (option == 2) { // 登录
            printf("用户名: ");
            fgets(username, MAX_NAME, stdin);
            username[strcspn(username, "\n")] = '\0';
            
            // 连接TCP
            if (connect(tcp_sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
                perror("TCP连接失败");
                close(udp_sock);
                close(tcp_sock);
                exit(1);
            }

            // 登录用户
            login_user(tcp_sock, username);
            if (strcmp(username, "") != 0) {
                logged_in = 1; // 设置已登录状态
            }
        } else if (option == 3) { // 退出程序
            printf("退出程序\n");
            break;
        } else {
            printf("无效选项。\n");
            continue; // 无效选项时重新提示
        }

        // 如果用户已登录
        if (logged_in) {
            // 关闭UDP套接字，后续使用TCP通信
            close(udp_sock);

            // 设置文件描述符集合
            fd_set read_fds;
            int max_fd = (tcp_sock > STDIN_FILENO) ? tcp_sock : STDIN_FILENO;

            printf("\n===== 聊天系统已启动 =====\n");
            print_menu();  // 先显示菜单

            while (1) {
                FD_ZERO(&read_fds);
                FD_SET(STDIN_FILENO, &read_fds);
                FD_SET(tcp_sock, &read_fds);

                if (select(max_fd + 1, &read_fds, NULL, NULL, NULL) < 0) {
                    perror("select错误");
                    break;
                }

                // 处理用户输入
                if (FD_ISSET(STDIN_FILENO, &read_fds)) {
                    char input[10];
                    fgets(input, sizeof(input), stdin);
                    input[strcspn(input, "\n")] = '\0'; // 移除换行符

                    if (strcmp(input, "-help") == 0) {
                        print_menu(); // 显示菜单
                        continue; // 继续循环
                    }

                    int choice = atoi(input);

                    ChatMessage msg;
                    msg.type = 0;
                    strncpy(msg.sender, username, MAX_NAME);

                    switch (choice) {
                        case 1: // 获取在线用户列表
                            msg.type = LIST_USERS;
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            break;

                        case 2: // 获取群组列表
                            msg.type = LIST_GROUPS;
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            break;

                        case 3: { // 发送私信
                            char target[MAX_NAME], message[MAX_BUFFER];
                            printf("接收者用户名: ");
                            fgets(target, MAX_NAME, stdin);
                            target[strcspn(target, "\n")] = '\0';

                            printf("消息内容: ");
                            fgets(message, MAX_BUFFER, stdin);
                            message[strcspn(message, "\n")] = '\0';

                            msg.type = PRIVATE_MSG;
                            strncpy(msg.target, target, MAX_NAME);
                            strncpy(msg.message, message, MAX_BUFFER);
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            break;
                        }

                        case 4: { // 创建群组
                            char group[MAX_NAME];
                            printf("新群组名称: ");
                            fgets(group, MAX_NAME, stdin);
                            group[strcspn(group, "\n")] = '\0';

                            msg.type = CREATE_GROUP;
                            strncpy(msg.target, group, MAX_NAME);
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            break;
                        }

                        case 5: { // 加入群组
                            char group[MAX_NAME];
                            printf("要加入的群组: ");
                            fgets(group, MAX_NAME, stdin);
                            group[strcspn(group, "\n")] = '\0';

                            msg.type = JOIN_GROUP;
                            strncpy(msg.target, group, MAX_NAME);
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            break;
                        }

                        case 6: { // 发送群组消息
                            char group[MAX_NAME], message[MAX_BUFFER];
                            printf("群组名称: ");
                            fgets(group, MAX_NAME, stdin);
                            group[strcspn(group, "\n")] = '\0';

                            printf("消息内容: ");
                            fgets(message, MAX_BUFFER, stdin);
                            message[strcspn(message, "\n")] = '\0';

                            msg.type = GROUP_MSG;
                            strncpy(msg.target, group, MAX_NAME);
                            strncpy(msg.message, message, MAX_BUFFER);
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            break;
                        }

                        case 7: { // 离开群组
                            char group[MAX_NAME];
                            printf("要离开的群组: ");
                            fgets(group, MAX_NAME, stdin);
                            group[strcspn(group, "\n")] = '\0';

                            msg.type = LEAVE_GROUP;
                            strncpy(msg.target, group, MAX_NAME);
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            break;
                        }

                        case 8: // 退出登录
                            msg.type = LOGOUT;
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            printf("已退出登录\n");
                            logged_in = 0; // 重置登录状态
                            // 重新进入选择注册或登录界面
                            close(tcp_sock); // 关闭TCP连接
                            tcp_sock = socket(AF_INET, SOCK_STREAM, 0); // 重新创建tcp_socket以便于后续的登录
                            break;

                        case 9: { // 广播消息
                            char message[MAX_BUFFER];
                            printf("广播消息内容: ");
                            fgets(message, MAX_BUFFER, stdin);
                            message[strcspn(message, "\n")] = '\0';

                            msg.type = BROADCAST;
                            strncpy(msg.message, message, MAX_BUFFER);
                            send(tcp_sock, &msg, sizeof(msg), 0);
                            break;
                        }

                        default:
                            printf("无效选项\n");
                            continue;
                    }
                }

                // 处理服务器消息
                if (FD_ISSET(tcp_sock, &read_fds)) {
                    ChatMessage recv_msg;
                    ssize_t n = recv(tcp_sock, &recv_msg, sizeof(recv_msg), 0);
                    if (n <= 0) {
                        printf("服务器连接已断开\n");
                        break;
                    }

                    switch (recv_msg.type) {
                        case PRIVATE_MSG:
                            printf("\n[私信] %s 对你说: %s\n", recv_msg.sender, recv_msg.message);
                            break;

                        case GROUP_MSG:
                            printf("\n[群组] %s@%s: %s\n", recv_msg.sender, recv_msg.target, recv_msg.message);
                            break;

                        case BROADCAST:
                            printf("\n[广播] %s: %s\n", recv_msg.sender, recv_msg.message);
                            break;

                        case LIST_USERS:
                            printf("\n=== 在线用户列表 ===\n%s\n", recv_msg.message);
                            break;

                        case LIST_GROUPS:
                            printf("\n=== 群组列表 ===\n%s\n", recv_msg.message);
                            break;

                        default:
                            printf("\n[系统通知] %s\n", recv_msg.message);
                    }
                }
            }
        }
    }

    close(tcp_sock); // 关闭TCP连接
    close(udp_sock); // 关闭UDP连接
    return 0;
}